package giteeai.impl;

import cc.git.liuyan.customeraiagent.core.bigmodel.*;
import cc.git.liuyan.customeraiagent.core.embeddingmodel.EmbeddingModelOutputData;
import cc.git.liuyan.customeraiagent.core.suppertool.KnowledgeSupperTool;
import cc.git.liuyan.customeraiagent.core.suppertool.Text2SqlSupperTool;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public class SimpleBigModelChatTool implements BigModelChatTool {
    Map<String, List<BigModelChatData>> datasMap = new ConcurrentHashMap<>();

    @Override
    public BigModelChatOutputData customReplay(BigModelChatInputData inputData) {
        if (inputData.getContent().equals("你是由谁开发的")) {
            return new BigModelChatOutputData("我是由刘颜开发的CustoemrAiAgent智能客服");
        } else {
            return null;
        }
    }

    @Override
    public List<BigModelChatData> loadChatData(BigModelChatInputData inputData) {
        String mapKey = inputData.getUserId() + "@#@" + inputData.getAgentId() + "@#@" + "ChatData";
        return datasMap.getOrDefault(mapKey, null);
    }

    @Override
    public void saveChatData(BigModelChatInputData inputData, BigModelChatData chatData) {
        String mapKey = inputData.getUserId() + "@#@" + inputData.getAgentId() + "@#@" + "ChatData";
        if (datasMap.containsKey(mapKey)) {
            datasMap.computeIfPresent(mapKey, (k, v) -> {
                ArrayList<BigModelChatData> chatData1 = new ArrayList<>(v);
                chatData1.add(chatData);
                return chatData1;
            });
        } else {
            datasMap.put(mapKey, Collections.singletonList(chatData));
        }
    }

    @Override
    public void persistenceBigModelChatInputData(BigModelChatInputData inputData, String outputData) {
        return;
    }

    @Override
    public String getText2SqlFinalPrompt(BigModelChatInputData inputData) {
        String promptTemplate = "你是一名Mysql专家，现在需要阅读并理解下面的【数据库schema】描述，以及可能用到的【参考信息】，并运用Mysql知识生成sql语句回答【用户问题】。\n" +
                "【用户问题】\n" +
                "%s\n" +
                "\n" +
                "【数据库schema】\n" +
                "%s\n" +
                "\n" +
                "【参考信息】\n" +
                "%s\n" +
                "\n" +
                "【用户问题】\n" +
                "%s\n" +
                "\n" +
                "```sql";
        EmbeddingModelOutputData outputVectors = Text2SqlSupperTool.question2Vectors(inputData.getContent());
        List<String> schemas = Text2SqlSupperTool.relationSchemaSearch(outputVectors.getDenseVectors(), outputVectors.getSparseVectors());
        List<String> remarks = Text2SqlSupperTool.relationRemarkSearch(outputVectors.getDenseVectors(), outputVectors.getSparseVectors());
        return String.format(promptTemplate, inputData.getContent(), String.join("\n", schemas), String.join("\n", remarks), inputData.getContent());
    }

    @Override
    public String getChatFinalPrompt(BigModelChatInputData inputData) {
        if (inputData.getContent().contains("不使用知识库")) {
            return inputData.getContent();
        }
        EmbeddingModelOutputData outputVectors = KnowledgeSupperTool.question2Vectors(inputData.getContent());
        inputData.getUserId();
        inputData.getAgentId();
        //假设智能体挂载了N个知识库
        List<Long> knowledgeIds = Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L);
        List<String> knowledges = KnowledgeSupperTool.relationKnowledgeSearch(knowledgeIds, outputVectors.getDenseVectors(), outputVectors.getSparseVectors());
        //假设根据userId+agentId能找到如下模板
        return PromptUtil.format("**你要根据用户输入的问题：**\n%s\n\n**参考如下内容：**\n%s\n\n整理处理最终结果。", inputData.getContent(), String.join("\n", knowledges));
    }
}
